# 导入 pymilvus 连接模块
from pymilvus import connections, FieldSchema, DataType, CollectionSchema, Collection

# 连接到 Milvus 服务器（本地或远程）
# 使用默认连接别名，连接到本地主机的19530端口，数据库名为rensheng
connections.connect("default", host="42.194.164.248", port="19530", db_name="rensheng")

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128),
]

schema = CollectionSchema(fields=fields, description="Example collection schema")

# 检测集合是否存在，不存在则新建
collection_name = "example_collection"
try:
    # 尝试获取已存在的集合
    collection = Collection(collection_name)
    print(f"Collection '{collection_name}' created.")
except Exception as e:
    # 如果集合不存在，则创建新集合
    collection = Collection(collection_name, schema)
    print(f"Collection '{collection_name}' 创建成功.")

# 定义索引参数配置
index_params = {
    "metric_type": "L2",  # 距离度量方式：L2、IP、COSINE等
    "index_type": "IVF_FLAT",  # 索引类型：FLAT Index、IVF_FLAT（InVerted File with Flat倒排文件+暴力遍历）、HNSW等
    "params": {
        "nlist": 128
    },  # 索引参数配置，nlist表示聚类中心的数量 值越大召回率越高、速度越慢
}

# 为向量字段创建索引
collection.create_index(field_name="embedding", index_params=index_params)
print(f"Index created for field 'embedding' in collection '{collection_name}'.")

# 将集合加载到内存中以支持搜索操作
collection.load()
# 打印集合加载成功信息
print("集合已加载到内存")
